import asyncio
import json
from websockets.connection import State
import websockets

# WebSocket 服务器的 URL
WS_URL = "ws://127.0.0.1:2121/ws"  # 替换为实际的 WebSocket 服务器地址

# 用户名和密码 (示例)
USERNAME = "user01"
PASSWORD = "user_key01"

# 重连间隔时间（秒）
RECONNECT_INTERVAL = 5
# 心跳间隔时间（秒），根据服务器要求调整
PING_INTERVAL = 20
# 新增：Ping超时时间（秒）
PING_TIMEOUT = 10

async def login(websocket):
    """向服务器发送登录信息"""
    login_data = {
        'user_code': USERNAME,
        'user_key': PASSWORD,
        'server_password': "myServerPassword"
    }
    await websocket.send(json.dumps(login_data))
    response = await websocket.recv()
    print(f"收到登录响应: {response}")
    return json.loads(response)

async def timer_ping(websocket, interval):
    """增强版心跳：检测Pong响应超时"""
    try:
        while True:
            await asyncio.sleep(interval)
            if websocket.state == State.OPEN:
                print("发送心跳 ping...")
                try:
                    # 发送Ping并等待Pong响应，超时则触发重连
                    await asyncio.wait_for(websocket.ping(), timeout=PING_TIMEOUT)
                    print("收到 pong 响应")
                except asyncio.TimeoutError:
                    # 主动抛出断开异常，触发重连
                    raise websockets.ConnectionClosedError(rcvd=None, sent=None)
            else:
                print("WebSocket 连接已关闭，停止发送 ping")
                break
    except Exception as e:
        print(f"发送 ping 出错: {e}")


async def receive_messages(websocket, loop):
    """接收来自服务器的消息"""
    message_count = 0
    try:
        async for message in websocket:
            message_count += 1
            msgData = json.loads(message)
            displayStr = msgData['message']['title']
		    print(f"收到消息 {message_count}: {displayStr}")
    except websockets.ConnectionClosed as e:
        print(f"连接已关闭: {e}")
        await reconnect(loop)

async def reconnect(loop):
    """尝试重新连接 WebSocket 服务器"""
    while True:
        try:
            print("尝试重新连接...")
            await main_loop(loop)
            break  # 如果成功连接，则退出循环
        except Exception as e:
            print(f"重连失败: {e}")
            await asyncio.sleep(RECONNECT_INTERVAL)  # 等待一段时间再尝试重连

async def main_loop(loop):
    """主函数：建立连接、登录并开始接收消息"""
    try:
        async with websockets.connect(WS_URL) as websocket:
            print("已连接到服务器")

            # 登录
            response = await login(websocket)
            if response.get('message') == '登录成功':
                print("登录成功")
                
                # 创建两个任务：接收消息 和 发送 ping
                receiver_task = asyncio.create_task(receive_messages(websocket, loop))
                ping_task = asyncio.create_task(timer_ping(websocket, PING_INTERVAL))

                # 等待任意一个任务完成（比如连接断开）
                done, pending = await asyncio.wait(
                    [receiver_task, ping_task],
                    return_when=asyncio.FIRST_COMPLETED
                )

                # 取消未完成的任务
                for task in pending:
                    task.cancel()

                # 如果不是正常关闭，触发重连
                raise websockets.ConnectionClosedError(rcvd=None, sent=None)

            else:
                print("登录失败")
    except websockets.ConnectionClosedOK:
        print("连接正常关闭")
    except Exception as e:
        print(f"主循环异常: {e}")
        await reconnect(loop)

async def start_client():
    """启动 WebSocket 客户端"""
    loop = asyncio.get_running_loop()
    while True:
        try:
            await main_loop(loop)
        except Exception as e:
            print(f"发生错误: {e}")
            await asyncio.sleep(RECONNECT_INTERVAL)

def graceful_exit():
    print("程序正在退出...")
    print("退出成功！再见！")

if __name__ == "__main__":
    try:
        asyncio.run(start_client())
    except KeyboardInterrupt:
        graceful_exit()
